#!/usr/bin/env python3
"""Run the 1D Helmholtz locELM reproduction experiment."""
from __future__ import annotations

import argparse
from time import perf_counter

import numpy as np

from reproduct_paper import LocELMSolver1D
from reproduct_paper.problems import helmholtz_1d_reference


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--subdomains", type=int, default=4, help="number of sub-domains")
    parser.add_argument(
        "--collocation",
        type=int,
        default=100,
        help="collocation points per sub-domain",
    )
    parser.add_argument(
        "--width",
        type=int,
        default=100,
        help="hidden units (training parameters) per sub-domain",
    )
    parser.add_argument(
        "--rand-mag",
        type=float,
        default=3.0,
        help="maximum magnitude of random hidden parameters",
    )
    parser.add_argument("--seed", type=int, default=1, help="seed for RNG")
    parser.add_argument(
        "--samples",
        type=int,
        default=2000,
        help="evaluation samples for error metrics",
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    config = helmholtz_1d_reference()
    rng = np.random.default_rng(args.seed)

    solver = LocELMSolver1D(
        domain=config["domain"],
        num_subdomains=args.subdomains,
        collocation_per_subdomain=args.collocation,
        hidden_width=args.width,
        rand_magnitude=args.rand_mag,
        rhs=lambda x: np.asarray(config["rhs"](x)).item(),
        boundary=lambda x: np.asarray(config["boundary"](x)).item(),
        lam=config["lambda"],
        rng=rng,
    )

    start = perf_counter()
    solver.solve()
    elapsed = perf_counter() - start

    # Evaluate on dense grid
    x_eval = np.linspace(config["domain"][0], config["domain"][1], args.samples)
    numerical = solver.evaluate(x_eval)
    exact = config["exact"](x_eval)
    error = numerical - exact
    max_error = np.max(np.abs(error))
    rms_error = np.sqrt(np.mean(error ** 2))

    print("locELM 1D Helmholtz experiment")
    print(f"  sub-domains:          {args.subdomains}")
    print(f"  collocation/subdom:   {args.collocation}")
    print(f"  hidden width/subdom:  {args.width}")
    print(f"  random magnitude:     {args.rand_mag}")
    print(f"  solve time (s):       {elapsed:.3f}")
    print(f"  max abs error:        {max_error:.3e}")
    print(f"  RMS error:            {rms_error:.3e}")


if __name__ == "__main__":
    main()
